# -*- coding: utf-8 -*-
import cv2
import numpy as np
import matplotlib.pyplot as plt


# 反相灰度图，将黑白阈值颠倒
def accessPiexl(img):
    # 获取图像的高度
    height = img.shape[0]
    # 获取图像的宽度
    width = img.shape[1]
    # 遍历图像的每个像素，将像素值反转（255 减去当前像素值）
    for i in range(height):
       for j in range(width):
           img[i][j] = 255 - img[i][j]
    # 返回反转后的图像
    return img


# 反相二值化图像
def accessBinary(img, threshold=128):
    # 调用 accessPiexl 函数反转图像的灰度值
    img = accessPiexl(img)
    # 定义一个 3x3 的全 1 核，用于膨胀操作
    kernel = np.ones((3, 3), np.uint8)
    # 对图像进行膨胀操作，迭代 1 次
    img = cv2.dilate(img, kernel, iterations=1)
    # 将图像中大于阈值的像素置为 0，小于等于阈值的像素保持不变
    _, img = cv2.threshold(img, threshold, 0, cv2.THRESH_TOZERO)
    # 返回处理后的图像
    return img


# 根据长向量找出顶点
def extractPeek(array_vals, min_vals=10, min_rect=20):
    # 存储提取的顶点对的列表
    extrackPoints = []
    # 起始点，初始化为 None
    startPoint = None
    # 结束点，初始化为 None
    endPoint = None
    # 遍历数组中的每个元素及索引
    for i, point in enumerate(array_vals):
        # 如果元素值大于阈值且起始点未设置，则设置起始点
        if point > min_vals and startPoint == None:
            startPoint = i
        # 如果元素值小于阈值且起始点已设置，则设置结束点
        elif point < min_vals and startPoint!= None:
            endPoint = i
        # 如果起始点和结束点都已设置，则将顶点对添加到列表中，并重置起始点和结束点
        if startPoint!= None and endPoint!= None:
            extrackPoints.append((startPoint, endPoint))
            startPoint = None
            endPoint = None
    # 剔除一些长度小于 min_rect 的矩形（噪点）
    for point in extrackPoints:
        if point[1] - point[0] < min_rect:
            extrackPoints.remove(point)
    # 返回提取的顶点对列表
    return extrackPoints


# 寻找边缘，返回边框的左上角和右下角（利用 cv2.findContours）
def findBorderContours(path, maxArea=50):
    # 以灰度模式读取图像
    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    # 对图像进行二值化和反相处理
    img = accessBinary(img)
    # 寻找轮廓，cv2.RETR_EXTERNAL 表示只检测最外层轮廓，cv2.CHAIN_APPROX_NONE 存储所有的轮廓点
    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    # 存储边框的列表
    borders = []
    # 遍历每个轮廓
    for contour in contours:
        # 计算轮廓的外接矩形
        x, y, w, h = cv2.boundingRect(contour)
        # 如果矩形面积大于 maxArea，则将其添加到边框列表中
        if w*h > maxArea:
            border = [(x, y), (x+w, y+h)]
            borders.append(border)
    # 返回边框列表
    return borders


# 显示结果及边框
def showResults(path, borders, results=None):
    # 读取原始图像
    img = cv2.imread(path)
    # 打印图像的形状
    print(img.shape)
    # 遍历每个边框
    for i, border in enumerate(borders):
        # 在图像上绘制矩形，边框颜色为红色
        cv2.rectangle(img, border[0], border[1], (0, 0, 255))
        # 如果有结果，则在边框左上角添加结果数字
        if results:
            cv2.putText(img, str(results[i]), border[0], cv2.FONT_HERSHEY_COMPLEX, 0.8, (0, 255, 0), 1)

        cv2.circle(img, border[0], 1, (0, 255, 0), 0)
    # 显示图像
    cv2.imshow('test', img)
    # 等待用户按键
    cv2.waitKey(0)


# 根据边框转换为 MNIST 格式
def transMNIST(path, borders, size=(28, 28)):
    # 创建一个存储转换后图像数据的数组，形状为 (边框数量, 28, 28, 1)
    imgData = np.zeros((len(borders), size[0], size[0], 1), dtype='uint8')
    # 以灰度模式读取图像
    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    # 对图像进行二值化和反相处理
    img = accessBinary(img)
    # 遍历每个边框
    for i, border in enumerate(borders):
        # 提取边框内的图像区域
        borderImg = img[border[0][1]:border[1][1], border[0][0]:border[1][0]]
        # 根据最大边缘拓展像素，使图像更接近正方形
        extendPiexl = (max(borderImg.shape) - min(borderImg.shape)) // 2
        # 对图像进行边界扩展，填充值为常数
        targetImg = cv2.copyMakeBorder(borderImg, 7, 7, extendPiexl + 7, extendPiexl + 7, cv2.BORDER_CONSTANT)
        # 调整图像大小为 28x28
        targetImg = cv2.resize(targetImg, size)
        # 为图像添加一个维度，以匹配神经网络的输入格式
        targetImg = np.expand_dims(targetImg, axis=-1)
        # 将处理后的图像存储在 imgData 中
        imgData[i] = targetImg
    # 返回转换后的图像数据
    return imgData


# 预测手写数字
def predict(modelpath, imgData):
    # 导入 Keras 的 models 模块
    from keras import models
    # 加载预训练的模型
    my_mnist_model = models.load_model(modelpath)
    # 打印模型的结构
    print(my_mnist_model.summary())
    # 将图像数据转换为浮点数并归一化到 0 到 1 的范围
    img = imgData.astype('float32') / 255
    # 使用模型进行预测
    results = my_mnist_model.predict(img)
    # 存储预测结果的列表
    result_number = []
    # 遍历预测结果，找出最大概率对应的数字
    for result in results:
        result_number.append(np.argmax(result))
    # 返回预测的数字列表
    return result_number


if __name__ == '__main__':
    # 图像文件的路径
    path = '2022212831.jpg'
    # 预训练模型的路径
    model = 'my_mnist_model.h5'
    # 查找图像的边框
    borders = findBorderContours(path)
    # 将边框内的图像转换为 MNIST 格式
    imgData = transMNIST(path, borders)
    # 预测边框内的手写数字
    results = predict(model, imgData)
    # 显示结果和边框
    showResults(path, borders, results)
